import numpy as np
import pandas as pd
from chapter_2_get import gen_stock_data_table

'''
海龟策略
在股价超过过去N个交易日股价最高点时买入，在股价低于过去N个交易日的股价最低点时卖出。
上述的若干个最高点和最低点会组成一个通道---'唐奇安通道'
'''

'''
使用过去N天的股价最高点和过去N天的股价最低点生成唐奇安通道。
一般N=20
'''


def turtle_strategy(symbol, start_date, end_date):
    # 创建一个turtle的数据表，使用原始数据的日期序列号
    stock_data = gen_stock_data_table(symbol, start_date, end_date)
    stock_002419 = stock_data['stock']
    print("-------------stock_002419_turtle--------------------")
    print(stock_002419)

    turtle = pd.DataFrame(index=stock_002419.index)
    print("-------------turtle--------------------")

    # 通道上界=过去20日内的最高价
    # 通道下界=过去20日内的最低价
    # 中轨道=0.5*（通道上界+通道下界）
    # 为什么要用shift(1)要向下移1？
    turtle['high'] = stock_002419['close'].shift(1).rolling(20).max()
    turtle['low'] = stock_002419['close'].shift(1).rolling(20).min()
    turtle['mid'] = (turtle['high'] + stock_002419['low']) / 2

    turtle['price'] = stock_002419['close']
    # 当股价突破下沿时卖出，发出卖出信号，反之买入
    # 代码实现是简化的买入卖出策略，并不是实际的海龟策略算法的买入卖出时机。
    turtle['buy'] = stock_002419['close'] > turtle['high']
    turtle['sell'] = stock_002419['close'] < turtle['low']
    print(turtle)

    # 订单初始状态为0
    turtle['order'] = 0
    # 前一天的仓位
    turtle['preposition'] = 0
    # 当前仓位
    turtle['position'] = 0

    # 初始仓位为0
    position = 0

    # 设置循环，便利turtle数据表
    for k in range(len(turtle)):
        # 当买入信号为true时，且仓位为0时，下单买入一手
        if turtle.buy[k] and position == 0:
            turtle.preposition.values[k] = position  # 保存前一天仓位
            turtle.order.values[k] = 1
            position = 1
            turtle.position.values[k] = position  # 修改当前仓位

        elif turtle.sell[k] and position > 0:
            turtle.preposition.values[k] = position  # 保存前一天仓位
            turtle.order.values[k] = -1
            position = 0
            turtle.position.values[k] = position  # 修改当前仓位

    return turtle


# turtle_strategy('002419', '20210101', '20221231')